#!/usr/bin/env python3

import cv2 as cv
import numpy as np

img = cv.imread('../pic/digits.png')
gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
print(gray.shape)

# 现在我们将图像分割为5000个单元格，每个单元格为20x20    vsplit分割成50行 hsplit分割每行为100列 indices_or_sections必须被整除
cells = [np.hsplit(row, indices_or_sections=100) for row in np.vsplit(gray, indices_or_sections=50)]
# 使其成为一个Numpy数组。x的大小将是（50,100,20,20）四维数组
x = np.array(cells)

# 现在我们准备train_data和test_data
# 对于每个数字，我们将其展平为一行，每行 400 个像素
train = x[:, :50].reshape(-1, 400).astype(np.float32)  # Size = (2500,400)
test = x[:, 50:100].reshape(-1, 400).astype(np.float32)  # Size = (2500,400)
# 为训练和测试数据创建标签
k = np.arange(10)
train_labels = np.repeat(k, 250)[:, np.newaxis]
# 二维数组 2500，1
test_labels = train_labels.copy()
# 初始化kNN
knn = cv.ml.KNearest_create()
# 使用训练数据和标签训练模型
knn.train(train, cv.ml.ROW_SAMPLE, train_labels)
ret, result, neighbours, dist = knn.findNearest(test, k=5)
# 检查分类的准确性
# 将结果与test_labels进行比较，并检查哪个错误
matches = result == test_labels
correct = np.count_nonzero(matches)
accuracy = correct * 100.0 / result.size
print(accuracy)

# 读取和保存数据
np.savez('knn_data.npz', train=train, train_labels=train_labels)
# Now load the data
with np.load('knn_data.npz') as data:
    print(data.files)
    train = data['train']
    train_labels = data['train_labels']
